#!/usr/bin/env python3
"""
过滤mitmproxy流量记录，只保留URL中包含anyrouter的请求信息
"""

import re
import sys

def filter_anyrouter_traffic(input_file, output_file):
    """
    过滤流量记录，只保留包含anyrouter的请求
    """
    try:
        with open(input_file, 'r', encoding='utf-8') as f:
            content = f.read()
    except FileNotFoundError:
        print(f"错误: 找不到文件 {input_file}")
        return False
    except Exception as e:
        print(f"错误: 读取文件失败 - {e}")
        return False

    # 按交互分割内容
    interactions = content.split('================================================================================')
    
    filtered_interactions = []
    anyrouter_count = 0
    
    # 保留文件头部信息
    if interactions:
        header = interactions[0]
        filtered_interactions.append(header)
    
    # 处理每个交互记录
    for i, interaction in enumerate(interactions[1:], 1):
        if not interaction.strip():
            continue
            
        # 提取URL并检查是否包含anyrouter
        url_match = re.search(r'URL: (https?://[^\s]+)', interaction)
        if url_match and 'anyrouter' in url_match.group(1).lower():
            filtered_interactions.append('================================================================================')
            filtered_interactions.append(interaction)
            anyrouter_count += 1
            print(f"找到anyrouter请求 #{anyrouter_count}: {url_match.group(1)}")
    
    # 写入过滤后的内容
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            # 更新头部统计信息
            filtered_content = ''.join(filtered_interactions)
            
            # 替换原始统计信息
            filtered_content = re.sub(
                r'去重后记录数: \d+',
                f'包含anyrouter的记录数: {anyrouter_count}',
                filtered_content
            )
            filtered_content = re.sub(
                r'输出路径: [^\n]+',
                f'输出路径: {output_file}',
                filtered_content
            )
            
            f.write(filtered_content)
        
        print(f"\n过滤完成!")
        print(f"原始记录: 从输入文件读取")
        print(f"anyrouter记录: {anyrouter_count}")
        print(f"输出文件: {output_file}")
        
        return True
        
    except Exception as e:
        print(f"错误: 写入文件失败 - {e}")
        return False

def main():
    input_file = 'mitmproxy_traffic_claude_filtered.txt'
    output_file = 'anyrouter_traffic_filtered.txt'
    
    if len(sys.argv) > 1:
        input_file = sys.argv[1]
    if len(sys.argv) > 2:
        output_file = sys.argv[2]
    
    print(f"输入文件: {input_file}")
    print(f"输出文件: {output_file}")
    print("开始过滤anyrouter流量记录...")
    
    success = filter_anyrouter_traffic(input_file, output_file)
    
    if success:
        print("\n✅ 过滤成功完成!")
    else:
        print("\n❌ 过滤失败!")
        sys.exit(1)

if __name__ == '__main__':
    main()